Read a mapped upstream's XCom sequence in chunks instead of one item per request - #73790

…ts type

_execute_task asked twice whether the task instance was an
IndexedTaskInstance: to skip exporting AIRFLOW_CTX_* to os.environ and to
leave the execution timeout to the iterated task. Each class now answers
for itself: RuntimeTaskInstance.export_context_to_env() and
enforces_execution_timeout(), which IndexedTaskInstance overrides, so the
runner no longer needs to know the subclass. Boy scout rule on code this
change added.

Measured with a CodSpeed simulation (valgrind), base vs head, median of
three rounds, instructions:

  control (checkpoint read)            57,690 ->     57,813  +0.2%
  plain task through _execute_task    310,483 ->    313,530  +1.0% (rounds spread 3%)
  ten sync items                   74,464,963 -> 74,650,754  +0.25%

In isolation the method call costs the same as a failing isinstance on a
plain task instance (76 vs 80 ns, pydantic models go through
ABCMeta.__instancecheck__) and about 60 ns more per call on an indexed one.

…its neighbours

SubOperatorRegister was the last Sub* name next to IndexedTaskInstance,
IndexedTaskRunner, IndexedTaskState and IndexedTaskOutcomes; it becomes
IndexedOperatorRegister. The docstrings of IndexedTaskRunner say indexed task
instead of sub-task as well.

…or an item

The classes say IndexedTaskInstance, IndexedTaskRunner, IndexedTaskState and
IndexedOperatorRegister, while the functions, messages, docstrings, comments,
docs page and test names around them called the same thing a sub-task, a
sub-operator or an item. An item now only means the input value an index runs
with; iteration stays, as the feature's own word.

- keeps_failure_to_its_item, _run_item, _run_task, _report_item and
  _create_task become keeps_failure_to_its_indexed_task, _run_index,
  _run_indexed_task, _report_indexed_task and _create_indexed_task
- the errors name the indexed task ("Indexed task x[i] ...", "N of M
  indexed tasks ran")
- the docs page defines the iterated task and the indexed task

…pstream provides

The partial kwargs holding an XComArg are the same for every index, so they
are found once per iterated task and only they are walked for the
fingerprint; a task without one walks nothing per index. A success or skip
report without a callback only sets the state and end_date, which makes no
SDK call, so it runs on the loop instead of in a worker thread.

CodSpeed simulation over 20 indexed tasks: -4.1% instructions for sync
operators, -3.6% for async ones.

…hen they differ

Unmapping an index that iterates params resolved its values into the Param
objects a shallow copy shares with the mapped operator: iterate() unmaps every
index in one process, so one index's params became the next one's. They are
copied deeply when the index brings its own.

An index's own params also never reached execute: the run context was cloned
from the parent's again. They are processed once, when the built operator's
params differ from the iterated task's, and both the render and the run
context read them; any other index keeps the parent's processed params, which
the clone copies, instead of processing them again.

CodSpeed simulation over 20 indexed tasks: -2.2% instructions for sync
operators, -2.0% for async ones.

…edence

Only ImportError and TypeError were taken as a refusal: a pydantic result whose
model gained a required field since the checkpoint was written raises a
ValidationError, a ValueError, and the indexed task that had succeeded failed
every retry and every clear until the checkpoint expired. The checkpoint is
already read, so nothing transient can fail while restoring it, and any
exception is now taken as a refusal.

A partial kwarg from an iterated upstream renders to an XComIterable, whose
serialized form keeps a handle and no item: clearing that upstream with this
task and getting other items of the same length left the fingerprint as it
was, and the indexed tasks that had succeeded replayed results computed from
the old input. A mapped upstream's LazyXComSequence was fetched whole by its
serializer for every index.

Both are now digested by their items, read one at a time, once per run: the
first indexed task to ask takes the digest and the others reuse it. CodSpeed
simulation over 20 indexed tasks with a partial kwarg from an upstream of 100
items: -34.6% instructions and 2000 -> 100 item reads for a mapped upstream,
+3.4% and 0 -> 100 reads for an iterated one; flat without an upstream.

Only an XComIterable or a LazyXComSequence in a partial kwarg was digested
by its items. The sequences .map(), .zip() and .concat() wrap around an
upstream took the plain serde path: a _MapResult could not be serialized,
which left the whole index without a fingerprint, so its checkpoint was
replayed whatever had changed, the iterated input included; a _ZipResult or
_ConcatResult around an iterated upstream was digested by its handle, which
stays the same when the upstream is cleared and returns other items.

Every lazy sequence, anything but a str, bytes, list or tuple, is now
digested by its items. A value whose items cannot all be read gets no digest
rather than failing every indexed task: a .map() callable may raise for an
item the indexed tasks never read.

The digest of a lazy upstream in a partial kwarg reads every item, one
supervisor round trip each, before the first indexed task starts. A kill
landing meanwhile stopped the indexed tasks waiting for the digest, but not
the read itself, which went on to the last item and held the process until
the supervisor escalated.

The digest now asks the iteration's stop flag before every read, as the
executor does before every pull, and leaves the value without a digest once
it is set: no indexed task runs after the kill to use it. CodSpeed
simulation of the hot-path benchmarks: within 0.4% of the previous head.

…ts task state

serde writes a set's members in iteration order, which for strings follows the
hash seed, and sort_keys only orders dict keys: the same input fingerprinted
differently on another worker, and a retry there ran every succeeded index
again. The members of every set are now put in a fixed order before the
digest, in the input fingerprint and in the digest of a lazy upstream's items.
A value without a set is serialized exactly as before, so the fingerprints of
earlier checkpoints still match.

…ndex

Every indexed task serialized the partial kwargs an upstream provides again for
its fingerprint, although they are the same for every index: a lookup list or a
config from an upstream behind .partial() cost its whole serialization once per
index. They are now made canonical once per run, by the first indexed task that
asks, and each index only adds its own iterated input to the digest. Without a
partial kwarg from an upstream the fingerprint is the same as before.

CodSpeed simulation, 20 indexed tasks with a partial kwarg from a plain upstream
list of 1,000 rows: 1,258M -> 188M instructions (-85%), peak memory 2.1-2.6 MiB
-> 0.95 MiB; the cases without such a kwarg are unchanged within noise.

…d loop is not kept

An asyncio.Lock that has had a waiter keeps a reference to its loop, so as the
value of the decoder's per-loop WeakKeyDictionary it kept its own key alive:
a sync indexed task that runs asyncio.run() per index with overlapping calls
kept one closed loop per index until the task ended. A lock no call waited for
holds no such reference and is still reused across calls; one that a call
waited for is dropped when the loop's last call leaves.

The per-loop bookkeeping moves out of asend() into _LoopQueues, which owns the
weak-keyed map, and _LoopQueue, an async context manager per loop that counts
its calls, notes contention and leaves the map when it must.

CodSpeed simulation of the comms benchmark: one asend at a time +1.0%, 50
concurrent +0.4%, send unchanged; the wall-clock wait of asyncio.to_thread work
next to 50 calls in flight is unchanged (0.24 -> 0.26 ms median).

…gest

One item of a lazy upstream that could not be read or serialized, such as a
.map() callable raising for an element no indexed task reads, left the whole
value without a digest on every attempt, so the fingerprints still matched
after the upstream was cleared and produced other values. The items are now
read by position and an unreadable one is digested as a fixed marker naming its
exception's type, so the other items still count; a warning names the path, as
the indexed tasks then run again, or not, for a reason found in that upstream.

…ays its own

_get_unmap_kwargs deep-copied params only when an index iterated them. Otherwise
copy.copy kept the same Param objects, which BaseOperator.__init__ wraps again
without copying and ParamsDict.__setitem__ resolves a new value into, so an
operator setting a param in its __init__ changed the iterated task's params too
and every index read the parent's value, and a param one index's execute set
reached its siblings. The params are now always copied deeply.

With its own copy, an index whose operator set a param differs from the
iterated task again, so the shortcut that processes only such an index's params
stays: processing every index's params measured +1.7% (sync) and +2.9% (async)
instructions in CodSpeed simulation over 20 indexed tasks.

…s own

IterableOperator had grown to 33 methods doing six jobs; eight of them, and two
caches with their locks in IterationState, only served the fingerprint that
tells whether an indexed task's checkpoint still applies. InputFingerprint now
owns that question for one run: the partial kwargs that come from an upstream,
the walk of declared against rendered values, the digests of lazy upstream
values, the canonical JSON and both per-run caches. IterationState hands out the
run's instance, and _create_indexed_task asks it for the fingerprint.

Every test case is kept, with its call target moved; the one branch the old
code left untested (an upstream input serde cannot serialize) now has a test.
CodSpeed simulation shows every hot-path case within noise.

…vel helpers

InputFingerprint held two caches with their locks, the digests of lazy upstream
values and the canonical text of the partial kwargs from an upstream, and a
static digest() that took that text as an argument although it was the
instance's own state. _LazyDigests and _UpstreamInputs now own one cache each,
the canonical JSON is a module function the three share, and the digest is part
of InputFingerprint.of().

Every test case is kept with its call target moved; the module stays fully
covered, and CodSpeed simulation shows every hot-path case within noise.

…erator stays a Dag node

The review rounds of Task Iteration each fixed a real problem (kill paths,
checkpoint ordering, thread affinity, params, fingerprints), and each fix landed
where the code already was: as another method, parameter or flag on
IterableOperator. By the approval round the class was 1,033 lines with 33
methods. Every method was sound on its own; the class no longer had one
responsibility. The input fingerprint moved out first (two commits before this
one). This commit finishes the split, with no change in behaviour.

What was wrong, and the principle each fix applies:

- One run's executor, context, template environment, checkpoint flags and
  outcomes were passed from method to method through six methods
  (_run_index took seven parameters). Values that travel together are an
  object that is missing (Fowler, Introduce Parameter Object): IterationRun,
  one per execute(), holds them.
- One indexed task's task and runner travelled the same way through its
  steps (replay, publish, record skip, record failure, report). They are an
  object as well (Fowler, Extract Class): IndexedTaskLifecycle, one per index.
- Building an indexed task (unmap, params, render, fingerprint) used fields
  nothing else in the run touched: IndexedTaskUnmapper, as unmapping and
  rendering go together in MappedOperator.render_template_fields.
- IterationState held two jobs behind two locks: what a kill must reach,
  which exists before the run, and the run's own data. It keeps the kill
  bookkeeping and kills the indexed operators itself; the run's data moves to
  IterationRun (single responsibility, Martin's SRP).
- Static and class methods that used no state of their class were moved to
  where their data lives (Fowler, Move Function): _report_indexed_task to
  IndexedTaskLifecycle, whose executor it took; _retryable and _fail_fast_for,
  the runner's rules for an exception, to module functions;
  _refuse_operators_that_skip_downstream inlined next to the guard of the same
  kind in __init__; axcom_push, a one-line wrapper, to the lifecycle.
- Checkpoints also answers whether an index's checkpoint may be trusted, so
  trust_checkpoints and since stop travelling with each indexed task.

The new classes are collaborators IterableOperator holds or makes, never bases
it inherits from (composition over inheritance): it subclasses BaseOperator
because the Dag needs one, and delegates the run. It comes last in the module,
after everything it uses.

IterableOperator: 1,033 -> 399 lines (about 160 of them the class docstring),
33 -> 12 methods. No class in the module has more than 12 methods. ruff's
opt-in design rules (PLR0904, PLR0913, PLR0917, PLR6301), which flagged
_run_index, _run_indexed_task, _publish and axcom_push before, pass.

Tests: every test function of test_iterableoperator.py is kept under its name
(one moved from TestIterationState to TestIterationRun with the length it
tests); only call targets changed, through a run_of() helper that builds a run
as execute() does. Two tests are new: IterationState.kill reaching every
operator whatever one raises, and the unmapper's do_xcom_push. Line and branch
coverage of the module are unchanged (17 and 4 missed, base 18 and 4).

CodSpeed simulation, base vs head, two rounds each, plus peak memory:
| Benchmark      | base 9726e07   | head              | mean   |
|----------------|-------------------|-------------------|--------|
| sync items     | 156.42M / 154.50M | 152.44M / 153.99M | -1.4%  |
| async items    | 152.87M / 152.56M | 152.38M / 152.14M | -0.3%  |
| iterated upstr | 106.81M / 107.18M | 106.94M / 106.61M | -0.2%  |
| mapped upstr   | 107.53M / 105.84M | 106.51M / 105.40M | -0.7%  |
| plain upstr    | 189.08M / 187.85M | 188.34M / 186.83M | -0.5%  |
| control        |   0.517M / 0.519M |   0.518M / 0.519M |  0     |
| peak memory    | 1.02 / 0.86 MiB   | 0.95 / 0.86 MiB   |        |

Every case is within its round-to-round spread: the split costs nothing
measurable, in time or memory.

…ide the loop

Nothing tested .iterate() together with loops at parse time, nor what an
iterated task's input holds when it reads a loop task:

- inside the loop, over a task of the same loop, it validates and its input is
  the current pass's value;
- outside the loop (or in another loop), it validates too, waits for the loop's
  gate, and its input is a lazy sequence of every pass's value, as for any task
  that reads a loop task from outside it. .expand() over the same task is
  refused at parse time; these tests pin today's .iterate() behaviour until
  that is decided;
- each pass is a task instance of its own and writes its indexed tasks'
  checkpoints under it, so a later pass never replays an earlier one.

…) is

The loop check only looked at mapped tasks. An iterated task is not mapped,
so .iterate() over a task of a loop passed from outside that loop, where
.expand() over the same task is refused. It then read every pass's value
and ran one index per pass instead of one per item, without any warning.

Every operator now says which upstream tasks it runs once per item of:
iter_expansion_sources(), empty by default. A mapped task returns its mapped
dependencies; an iterated task the upstream tasks of its input, which are not
reported as mapped dependencies, as the iterated task is not mapped and the
scheduler must not count those upstream items.

The check starts from each loop's gate: a task outside a loop that reads
one of its tasks depends on the gate instead (get_edge_source), so only the
gate's downstream tasks can fan out over a task of the loop. It no longer
walks every mapped task of the Dag. CodSpeed simulation, the check and
DAG.validate() on three Dag shapes, two rounds each:

| Benchmark                  | before              | after             |
|----------------------------|---------------------|-------------------|
| check, 1000 plain, no loop |   4.15M /   4.17M   | 2.97M / 2.97M -29% |
| check, 1000 plain + a loop |   4.30M /   4.32M   | 3.13M / 3.12M -28% |
| check, 999 mixed + a loop  |  12.43M /  12.37M   | 3.10M / 3.12M -75% |
| validate(), 999 mixed      |  13.22M /  13.14M   | 3.87M / 3.86M -71% |

The .iterate() cases from outside a loop and from another loop move into
the rejecting test; the two tests that pinned the old behaviour go.

…per request

LazyXComSequence fetched one item per GetXComSequenceItem request: iterating a
mapped upstream's results, or an iterated task reading them by index, cost one
supervisor round trip per item, which is the dominant cost for large inputs.

Reads now come in chunks. An index outside the slice held fetches chunk_size
items starting there with one GetXComSequenceSlice request, the message the
slice path already used, and consecutive reads are served from it. Sequential
access costs one request per chunk instead of one per item, both through
__getitem__/__iter__ (send) and aget/__aiter__ (asend), and at most one chunk
is held. A jump backwards fetches again from there.

The chunk size is the new [core] xcom_sequence_chunk_size option, default 32,
documented in config.yml and in the iteration docs; a LazyXComSequence can
also be given one explicitly.

A negative index no longer needs its own per-item request: it is mapped from the
end through the cached length, one GetXComCount at most. The client no longer
sends GetXComSequenceItem at all; the request handler keeps serving it.